package org.zhuzx.util.spring;

import java.util.ArrayList;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import javax.sql.DataSource;

import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.ResultSetExtractor;

/**
 * JdbcTemplate的扩展，自动根据sql性质切换数据源。
 * @author zhuzx 2018年08月23日
 */
public class SwitchableJdbcTemplate extends JdbcTemplate {
	
	private static ThreadLocal<Boolean> needSwitch = new ThreadLocal<Boolean>() {
		protected Boolean initialValue() {
			return false;
		}
	};
	
	private DataSource masterDataSource;
	private DataSource slaveDataSource;
	private boolean useRadicalMode = false;
	
	public SwitchableJdbcTemplate(DataSource masterDataSource, DataSource slaveDataSource) {
		super(masterDataSource);
		this.slaveDataSource = slaveDataSource;
	}
	
	public SwitchableJdbcTemplate(DataSource masterDataSource, DataSource slaveDataSource,
			boolean useRadicalMode) {
		this(masterDataSource, slaveDataSource);
		this.useRadicalMode = useRadicalMode;
	}
	
	public void setDataSource(DataSource dataSource) {
		this.masterDataSource = dataSource;
	}
	
	public <T> T query(final String sql, final ResultSetExtractor<T> rse) throws DataAccessException {
		int level = ratingSql(sql);
		boolean flag = false;//用来避免不必要的set操作
		if (level > 0) {
			flag = true;
		} else if (level == 0) {
			flag = useRadicalMode;
		}
		if (flag) {
			needSwitch.set(true);
//			System.out.println(">>>>已切换到从库数据源");
		}
		T result = super.query(sql, rse);
		if (flag) {//还原
			needSwitch.set(false);
		}
		return result;
	}

	public DataSource getDataSource() {
		if (needSwitch.get()) {
			return this.slaveDataSource;
		} else {
			return this.masterDataSource;
		}
	}
	
	private static Pattern[] fastSqlFeatures;
	private static Pattern[] slowSqlFeatures;
	static {
		List<Pattern> fastList = new ArrayList<>();
		fastList.add(Pattern.compile(" \\w*_?id ?= ?\\d{1,}$"));
		fastSqlFeatures = fastList.toArray(new Pattern[fastList.size()]);
		
		String[] features = {
				"count", "sum", "max", "min", "avg", "group_concat", "group by"
		};
		slowSqlFeatures = new Pattern[features.length];
		for (int i=0; i<features.length; i++) {
			String feature = features[i];
			Pattern pat = Pattern.compile("\\b" + feature + "\\b");
			slowSqlFeatures[i] = pat;
		}
	}
	
	private static int ratingSql(String sql) {
		sql = sql.toLowerCase().trim().replaceAll("\\s+", " ");
		if (isFind(sql, fastSqlFeatures)) {
			return -1;
		} else if (isFind(sql, slowSqlFeatures)) {
			return 1;
		} else {
			return 0;
		}
	}
	
	private static boolean isFind(String str, Pattern[] patterns) {
		for (Pattern pat : patterns) {
			Matcher mat = pat.matcher(str);
			if (mat.find()) {
//				System.out.println("找到[" + mat.group() + "]：" + str);
				return true;
			}
		}
		return false;
	}
	
}